from .utils.mail import sendMail
from datetime import datetime
import re
from .utils.server import (
    convert_non_white_to_black,
    add_soft_outline_and_white_background,
)

import os
import random
import numpy as np
import torch
from PIL import Image, ImageOps, ImageSequence

# import os, glob, sys
# import comfy.utils
# import folder_paths
# from insightface.app.common import Face
# from safetensors.torch import save_file, safe_open

# models_dir = folder_paths.models_dir
# REACTOR_MODELS_PATH = os.path.join(models_dir, "reactor")
# FACE_MODELS_PATH = os.path.join(REACTOR_MODELS_PATH, "faces")

# from .utils.autonode import node_wrapper, get_node_names_mappings, validate, anytype

# fundamental_classes = []
# fundamental_node = node_wrapper(fundamental_classes)


class AlwaysEqualProxy(str):
    def __eq__(self, _):
        return True

    def __ne__(self, _):
        return False


class ExpressionNotice:
    """
    通过表达式，发送邮件提醒
    """

    def __init__(self):
        pass

    RETURN_TYPES = ("BOOLEAN",)
    RETURN_NAMES = ("boolean",)
    FUNCTION = "main"
    CATEGORY = "🤑acqui"

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "expression": (
                    "STRING",
                    {
                        "multiline": True,
                        "dynamicPrompts": False,
                    },
                ),
                "is_send_mail": (
                    "BOOLEAN",
                    {
                        "default": True,
                    },
                ),
                "remark": (
                    "STRING",
                    {
                        "multiline": False,  # True if you want the field to look like the one on the ClipTextEncode node
                        "default": "云端GPU! 4090",
                    },
                ),
            },
            "optional": {
                "a": ("INT,FLOAT,STRING,IMAGE,LATENT",),
                "b": ("INT,FLOAT,STRING,IMAGE,LATENT",),
                "c": ("INT,FLOAT,STRING,IMAGE,LATENT",),
            },
        }

    def main(
        self,
        expression,
        is_send_mail,
        remark,
        a=0,
        b=0,
        c=0,
    ):
        result = self.evaluate_expression(expression, int(a), int(b), int(c))
        print("表达式结果：", result)
        if is_send_mail & result:
            sendMail(remark)
        return (result,)

    def is_expression_valid(self, expr):
        # 使用正则表达式检查表达式是否只包含允许的字符
        allowed_chars = r"^[a-c0-9+\-*/().<> =]*$"
        if not re.match(allowed_chars, expr):
            raise ValueError("Invalid characters in the expression")
        return True

    def evaluate_expression(self, expression, a_value, b_value, c_value):
        # 检查表达式是否有效
        if not self.is_expression_valid(expression):
            return False

        # 将变量值替换到表达式中
        expression = expression.replace("a", str(a_value))
        expression = expression.replace("b", str(b_value))
        if c_value is not None:
            expression = expression.replace("c", str(c_value))

        try:
            # 使用eval()评估表达式
            result = eval(expression)
            return result
        except Exception as e:
            # 如果评估过程中发生错误，返回False
            print(f"Error evaluating expression: {e}")
            return False


class NumNoticeForceInput:
    """
    当数值到达指定数值时，发送邮件提醒
    """

    def __init__(self):
        pass

    RETURN_TYPES = ()
    FUNCTION = "detectionReminder"
    OUTPUT_NODE = True
    CATEGORY = "🤑acqui"

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "cureent_num": (
                    "INT",
                    {"default": 0, "step": 1, "forceInput": True},  # Slider's step
                ),
                "notice_num": (
                    "INT",
                    {
                        "default": 450,
                        "min": 1,  # Minimum value
                        "max": 4096,  # Maximum value
                        "forceInput": True,
                    },
                ),
                "remark": (
                    "STRING",
                    {
                        "multiline": False,  # True if you want the field to look like the one on the ClipTextEncode node
                        "default": "云端GPU! 4090",
                    },
                ),
            },
        }

    def detectionReminder(self, cureent_num, notice_num, remark):
        if cureent_num >= notice_num:
            sendMail(remark)
        return ()


class MaskFix:
    """
    遮罩优化
    """

    def __init__(self):
        pass

    FUNCTION = "main"
    CATEGORY = "🤑acqui"
    RETURN_NAMES = ("遮罩",)
    RETURN_TYPES = ("MASK",)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "mask": ("MASK",),
            },
        }

    def main(self, mask):
        return (convert_non_white_to_black(mask),)


class MaskOutline:
    """
    遮罩轮廓线
    """

    def __init__(self):
        pass

    FUNCTION = "main"
    CATEGORY = "🤑acqui"
    RETURN_NAMES = ("遮罩",)
    RETURN_TYPES = ("MASK",)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "mask": ("MASK",),
            },
        }

    def main(self, mask):
        return (add_soft_outline_and_white_background(mask),)


class TimeRecording:
    """
    记录当前时间
    """

    def __init__(self):
        pass

    FUNCTION = "main"
    CATEGORY = "🤑acqui"
    RETURN_NAMES = ("time",)
    RETURN_TYPES = ("STRING",)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "format": ("STRING", {"default": "记录时间 %Y-%m-%d %H:%M:%S"}),
            },
        }

    def main(self, format):
        return (datetime.now().strftime(format),)


# class LoadReactorFaceModel:
#     """
#     加载reactor模型
#     """

#     def __init__(self):
#         pass

#     FUNCTION = "main"
#     CATEGORY = "🤑acqui"
#     RETURN_NAMES = ("FACE_MODEL",)
#     RETURN_TYPES = ("FACE_MODEL",)

#     @classmethod
#     def INPUT_TYPES(s):
#         return {
#             "required": {
#                 "face_model_name": (
#                     "STRING",
#                     {
#                         "multiline": False,  # True if you want the field to look like the one on the ClipTextEncode node
#                         "default": "",
#                     },
#                 ),
#             },
#         }

#     def main(self, face_model_name):
#         self.face_model_name = face_model_name
#         self.face_models_path = FACE_MODELS_PATH
#         if self.face_model_name != "none":
#             face_model_path = os.path.join(self.face_models_path, self.face_model_name)
#             face = {}
#             with safe_open(face_model_path, framework="pt") as f:
#                 for k in f.keys():
#                     face[k] = f.get_tensor(k).numpy()
#             out = Face(face)
#         else:
#             out = None
#         return (out,)


class StrToColor:
    """
    字符串到COLOR类型
    """

    def __init__(self):
        pass

    FUNCTION = "main"
    CATEGORY = "🤑acqui"
    RETURN_NAMES = ("COLOR",)
    RETURN_TYPES = ("COLOR",)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "color_str": (
                    "STRING",
                    {
                        "multiline": False,  # True if you want the field to look like the one on the ClipTextEncode node
                        "default": "",
                    },
                ),
            },
        }

    def main(self, color_str):
        return (color_str,)


class StrToBoolean:
    """
    任意转布尔类型
    """

    def __init__(self):
        pass

    FUNCTION = "main"
    CATEGORY = "🤑acqui"
    RETURN_NAMES = ("BOOLEAN",)
    RETURN_TYPES = ("BOOLEAN",)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "str": (AlwaysEqualProxy("*"),),
            },
        }

    def main(self, str):
        return (bool(str),)


# 搜索 Value not in list
# 使用下面相关节点需要将“execution.py”中的662行中的“errors.append(error)”注释掉 否则会报 xxx not in ["$$parma1", "$$parma2", "$$parma3", "$$parma4"]
class StringParameter:
    """
    字符串参数
    """

    def __init__(self):
        pass

    FUNCTION = "main"
    CATEGORY = "🤑acqui"
    RETURN_NAMES = ("STRING",)
    RETURN_TYPES = ("STRING",)

    PREDEFINED_VALUES = ["$$parma1", "$$parma2", "$$parma3", "$$parma4"]

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "key_name": (s.PREDEFINED_VALUES,),
            },
            "optional": {
                "key_val": (
                    "STRING",
                    {
                        "multiline": True,  # True if you want the field to look like the one on the ClipTextEncode node
                        "default": "",
                    },
                ),
            },
        }

    def main(self, key_name="$$parma1", key_val=""):
        if key_name.find("$$") < 0:
            return (key_name,)
        elif key_val != "":
            return (key_val,)
        return ""


class IntParameter:
    """
    数字参数
    """

    def __init__(self):
        pass

    FUNCTION = "main"
    CATEGORY = "🤑acqui"
    RETURN_NAMES = ("INT",)
    RETURN_TYPES = ("INT",)

    PREDEFINED_VALUES = ["$$seed"]

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "key_name": (s.PREDEFINED_VALUES,),
            },
            "optional": {
                "key_val": (
                    "STRING",
                    {
                        "multiline": False,  # True if you want the field to look like the one on the ClipTextEncode node
                        "default": "",
                    },
                ),
            },
        }

    def main(self, key_name="$$seed", key_val=""):
        try:
            if key_name.find("$$") < 0:
                return (int(key_name),)
            elif key_val != "":
                return (int(key_val),)
            return (0,)
        except Exception:
            return (key_name,)


class ReservedParameters:
    """
    预留参数
    """

    def __init__(self):
        pass

    FUNCTION = "main"
    CATEGORY = "🤑acqui"
    RETURN_NAMES = ("STRING",)
    RETURN_TYPES = ("STRING",)

    PREDEFINED_VALUES = [
        "$$userSourceImage1",
        "$$userSourceImage2",
        "$$referenceImage1",
        "$$referenceImage2",
    ]

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "key_name": (s.PREDEFINED_VALUES,),
            },
            "optional": {
                "key_val": (
                    "STRING",
                    {
                        "multiline": False,  # True if you want the field to look like the one on the ClipTextEncode node
                        "default": "",
                    },
                ),
            },
        }

    def main(self, key_name="$$userSourceImage1", key_val=""):
        if key_name.find("$$") < 0:
            return (key_name,)
        elif key_val != "":
            return (key_val,)
        return ""


class systemParameters:
    """
    系统参数
    """

    def __init__(self):
        pass

    FUNCTION = "main"
    CATEGORY = "🤑acqui"
    RETURN_NAMES = ("STRING",)
    RETURN_TYPES = ("STRING",)

    PREDEFINED_VALUES = [
        "$$systemParma1",
        "$$systemParma2",
        "$$systemParma3",
        "$$systemParma4",
    ]

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "key_name": (s.PREDEFINED_VALUES,),
            },
            "optional": {
                "key_val": (
                    "STRING",
                    {
                        "multiline": True,  # True if you want the field to look like the one on the ClipTextEncode node
                        "default": "",
                    },
                ),
            },
        }

    def main(self, key_name="$$systemParma1", key_val=""):
        if key_name.find("$$") < 0:
            return (key_name,)
        elif key_val != "":
            return (key_val,)
        return ""


#
class RandomImageFromDir:
    """
    从input中的指定路径中加载随机图像
    """

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {"folder_path": ("STRING", {"multiline": False, "default": ""})}
        }

    RETURN_TYPES = ("IMAGE", "MASK")
    FUNCTION = "load_image"
    CATEGORY = "🤑acqui"

    def load_image(self, folder_path):
        print(111111, os.getcwd())
        currentPath = os.getcwd()
        if "/ComfyUI" in currentPath or "\ComfyUI\\" in currentPath:
            folder_path = os.path.join(currentPath, "input", folder_path)
        else:
            folder_path = os.path.join(currentPath, "ComfyUI", "input", folder_path)
        print(2222222, folder_path)

        files = os.listdir(folder_path)
        image_extensions = {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".ico", ".jfif"}
        images = [
            file
            for file in files
            if os.path.splitext(file)[1].lower() in image_extensions
        ]

        random_image = os.path.join(folder_path, random.choice(images))
        img = Image.open(random_image)
        output_images = []
        output_masks = []
        # this is from load_image node. Lazy but works :')
        for i in ImageSequence.Iterator(img):
            i = ImageOps.exif_transpose(i)
            if i.mode == "I":
                i = i.point(lambda i: i * (1 / 255))
            image = i.convert("RGB")
            image = np.array(image).astype(np.float32) / 255.0
            image = torch.from_numpy(image)[None,]
            if "A" in i.getbands():
                mask = np.array(i.getchannel("A")).astype(np.float32) / 255.0
                mask = 1.0 - torch.from_numpy(mask)
            else:
                mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
            output_images.append(image)
            output_masks.append(mask.unsqueeze(0))

        if len(output_images) > 1:
            output_image = torch.cat(output_images, dim=0)
            output_mask = torch.cat(output_masks, dim=0)
        else:
            output_image = output_images[0]
            output_mask = output_masks[0]

        return (output_image, output_mask)

    @classmethod
    def IS_CHANGED(s, image):
        return float("NaN")


class ThrowAnError:
    """
    抛出错误
    """

    def __init__(self):
        pass

    FUNCTION = "main"
    CATEGORY = "🤑acqui"
    RETURN_NAMES = ("BOOLEAN",)
    RETURN_TYPES = ("BOOLEAN",)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "boolData": (AlwaysEqualProxy("*"),),
            },
        }

    def main(self, boolData):
        if bool(boolData) == False:
            raise Exception("Error: {}".format("脚本执行错误"))
        return (bool(boolData),)


class MakeJson:
    """
    制作json
    """

    def __init__(self):
        pass

    FUNCTION = "main"
    CATEGORY = "🤑acqui"
    RETURN_NAMES = ("JSON",)
    RETURN_TYPES = ("JSON",)
    PREDEFINED_VALUES = [
        "NSFW",
    ]

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "key_name": (s.PREDEFINED_VALUES,),
                "val": (AlwaysEqualProxy("*"),),
            },
        }

    def main(self, key_name, val):
        obj = {key_name: val}
        return (obj,)


NODE_CLASS_MAPPINGS = {
    "ExpressionNotice": ExpressionNotice,
    "TimeRecording": TimeRecording,
    "NumNoticeForceInput": NumNoticeForceInput,
    "MaskFix": MaskFix,
    "MaskOutline": MaskOutline,
    #     "LoadReactorFaceModel": LoadReactorFaceModel,
    "StrToColor": StrToColor,
    "StrToBoolean": StrToBoolean,
    "StringParameter": StringParameter,
    "IntParameter": IntParameter,
    "ReservedParameters": ReservedParameters,
    "systemParameters": systemParameters,
    "ThrowAnError": ThrowAnError,
    "MakeJson": MakeJson,
    "RandomImageFromDir": RandomImageFromDir,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ExpressionNotice": "表达式邮箱通知",
    "NumNoticeForceInput": "表达式邮箱通知ForceInput",
    "TimeRecording": "日期format",
    "MaskFix": "遮罩优化",
    "MaskOutline": "遮罩轮廓线",
    #     "LoadReactorFaceModel": "Reactor加载面部模型",
    "StrToColor": "字符串到COLOR",
    "StrToBoolean": "字符串到Boolean",
    "StringParameter": "字符串参数",
    "IntParameter": "数字参数",
    "ReservedParameters": "预留参数(用于图像)",
    "systemParameters": "系统配置参数",
    "ThrowAnError": "抛出错误",
    "MakeJson": "制作JSON",
    "RandomImageFromDir": "从input下指定路径加载随机图像",
}

# WEB_DIRECTORY = "web"
